#include "vk_texture_array.h"
#include "stb_image.h"

namespace vkrender
{
    VulkanTextureArray::VulkanTextureArray(const std::shared_ptr<VulkanDeviceContext>& context,
        uint32_t width, uint32_t height, uint32_t max_textures)
    {
        this->context = context;
        this->width = width;
        this->height = height;
        this->max_textures = max_textures;
        this->num_textures = 0;

        // create image and memory for 
        init();
    }

    VulkanTextureArray::~VulkanTextureArray()
    {
        vkDestroyImage(context->device, imageArray, nullptr);
        vkFreeMemory(context->device, imageArrayMemory, nullptr);
    }

    bool VulkanTextureArray::addTexture(const std::string& texturePath)
    {
        if (num_textures >= max_textures)
            return false;

        int image_width, image_height, image_channels;
        stbi_uc* pixels = stbi_load(texturePath.c_str(), &image_width, 
            &image_height, &image_channels, STBI_rgb_alpha);

        // load image into buffer, copy(and resize) buffer data into texture array memory
        VkBuffer buffer;
        VkDeviceMemory bufferMemory;
        VkDeviceSize bufferSize = image_width * image_height * 4;
        context->createBuffer(bufferSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, buffer, bufferMemory);
        void* data;
        vkMapMemory(context->device, bufferMemory, 0, bufferSize, 0, &data);
        memcpy(data, pixels, static_cast<size_t>(bufferSize));
        vkUnmapMemory(context->device, bufferMemory);
        stbi_image_free(pixels);

        // copy into staging image
        VkImage stagingImage;
        VkDeviceMemory stagingImageMemory;
        VkImageCreateInfo imageInfo = defaultImageCreateInfo(image_width, image_height, 1);
        imageInfo.usage |= VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
        context->createImage(imageInfo, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, stagingImage, stagingImageMemory);
        context->transitionImageLayout(stagingImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
        context->copyBufferToImage(buffer, stagingImage, image_width, image_height);
        context->transitionImageLayout(stagingImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);

        // blit from staging memory into 
        context->transitionImageLayout(imageArray, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
        context->blitImage(stagingImage, image_width, image_height, imageArray, width, height, num_textures);
        context->transitionImageLayout(imageArray, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        vkDestroyBuffer(context->device, buffer, nullptr);
        vkFreeMemory(context->device, bufferMemory, nullptr);
        vkDestroyImage(context->device, stagingImage, nullptr);
        vkFreeMemory(context->device, stagingImageMemory, nullptr);

        num_textures += 1;
        return true;
    }

    void VulkanTextureArray::init()
    {
        VkDeviceSize size = width * height * 4 * max_textures;

        // create image/memory for texture array
        VkImageCreateInfo imageInfo = defaultImageCreateInfo(width, height, max_textures);
        context->createImage(imageInfo, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, this->imageArray, this->imageArrayMemory);
        // create imageview
        imageArrayView = context->createImageView(imageArray, VK_FORMAT_R8G8B8A8_SRGB, max_textures);

        // set image layout of all layers 
        context->transitionImageLayout(imageArray, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 0, max_textures);
    }

    VkImageCreateInfo VulkanTextureArray::defaultImageCreateInfo(uint32_t w, uint32_t h, uint32_t num_layers)
    {
        VkImageCreateInfo imageInfo = {};
        imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
        imageInfo.imageType = VK_IMAGE_TYPE_2D;
        imageInfo.extent.width = w;
        imageInfo.extent.height = h;
        imageInfo.extent.depth = 1;
        imageInfo.mipLevels = 1;
        imageInfo.arrayLayers = num_layers;
        imageInfo.format = VK_FORMAT_R8G8B8A8_SRGB;
        imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
        imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        imageInfo.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
        imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
        imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
        return imageInfo;
    }
}
